import itertools
import json
class ComboMutator:
    def __init__(self, mutator, config):
        self.mutator = mutator
        self.config = config
        self.rule_map = {
            'header': self.mutator.mutate_header,
            'filename': self.mutator.mutate_filename
        }
        
    def generate_mutations(self):
        if self.config['brute_mode']:
            yield from self._brute_force_combinations()
        else:
            yield from self._single_mutations()
    
    def _single_mutations(self):
        # 生成单变异请求
        yield {
            **self.mutator.original_request,
            'description': 'Original request'
        }
        
        if 'header' in self.config['allowed_rules']:
            mutated = self.mutator.mutate_header()
            mutated['description'] = 'Header mutation'
            yield mutated
            
        if 'filename' in self.config['allowed_rules']:
            mutated = self.mutator.mutate_filename()
            mutated['description'] = 'Filename mutation'
            yield mutated
    
    def _brute_force_combinations(self):
        # 生成组合变异请求
        rules = self.config['allowed_rules']
        for depth in range(self.config['min_depth'], self.config['max_depth']+1):
            for combo in itertools.permutations(rules, depth):
                mutated = self.mutator.original_request.copy()
                desc = []
                for rule in combo:
                    mutated = self.rule_map[rule](mutated)
                    desc.append(rule)
                mutated['description'] = '+'.join(desc)
                yield mutated
class ComboMutator:
    def __init__(self, mutator, config):
        # ...原有初始化...
        self.progress_file = "test_progress.json"
        self.current_state = {}
        
    def save_progress(self):
        with open(self.progress_file, 'w') as f:
            json.dump({
                'attempted': self.attempted,
                'successful': self.successful
            }, f)
            
    def resume_from(self, progress):
        self.attempted = progress['attempted']
        self.successful = progress['successful']